
from sqlalchemy.orm import Session
from src.db import engine
from src.db.table import CrawlerTask
from src.enum.task import CrawlerStatusEnum


def create(task: CrawlerTask):
    with Session(engine) as db:
        db.add(task)
        db.commit()
        db.refresh(task)
        return task

def get(crawler_id: str):
    with Session(engine) as db:
        task = db.query(CrawlerTask).filter(CrawlerTask.crawler_id == crawler_id).first()
        if not task:
            return None
        return task

def query_by_crawler_id(crawler_id: str):
    with Session(engine) as db:
        task = db.query(CrawlerTask).filter(CrawlerTask.crawler_id == crawler_id).first()
        if not task:
            return None
        return task

def query_all_by_crawler_status(crawler_status_list: list[CrawlerStatusEnum]):
    with Session(engine) as db:
        task_list = db.query(CrawlerTask).filter(CrawlerTask.crawler_status.in_(crawler_status_list)).all()
        return task_list

def update(task: CrawlerTask):
    with Session(engine) as db:
        task_db = db.query(CrawlerTask).filter(CrawlerTask.crawler_id == task.crawler_id).first()
        if not task_db:
            return None
        db.merge(task)
        db.commit()
        db.refresh(task_db)
        return task_db

def delete(crawler_id: str):
    with Session(engine) as db:
        task = db.query(CrawlerTask).filter(CrawlerTask.crawler_id == crawler_id).first()
        if task:
            db.delete(task)
            db.commit()